import sys, os
sys.path.append(r"E:\25spring\FYP\pymodules")
import numpy as np
import torch
from feature_extract.model_history import try_get_model, LSTM_v6
import matplotlib.pyplot as plt
import seaborn as sns
LSTM = LSTM_v6

import glob, json
import pandas as pd

MODEL_DIR = r"E:\25spring\FYP\pymodules\deep-models"
DATAD = r"C:\Users\songy\Desktop\data-0518"
# pths = glob.glob(f"{base_dir}/**/best_*.pth", recursive=True)

def get_dummy():
    static_left = []
    static_right = []
    move_left = []
    move_right = []
    for pos in ['l', 'r']:
        for freq in [10, 20, 30, 40, 50, 60]:
            dummy_static = [
                [0,0,0,0,0],
                [0,0,0,0,0],
                [0,0,0,0,0],
                [0,0,0,0,0],
                [0,0,0,freq,0] if pos == 'l' else [0,0,0,0,freq]
            ]
            dummy_move = [
                [0,0,0,0,0],
                [30,0,0,0,0],
                [60,0,0,0,0],
                [90,0,0,0,0],
                [120,0,0,freq,0] if pos == 'l' else [120,0,0,0,freq]
            ]
            dummy_static = torch.tensor(dummy_static, dtype=torch.float32).to("cuda")
            dummy_move = torch.tensor(dummy_move, dtype=torch.float32).to("cuda")
            if pos == 'l':
                static_left.append(dummy_static)
                move_left.append(dummy_move)
            else:
                static_right.append(dummy_static)
                move_right.append(dummy_move)
    return static_left, static_right, move_left, move_right

def step3(model, dummy):
    # take 3 steps
    sti = dummy[-1,3:]
    # print(f'input1: {dummy.shape}')
    feat1 = model(dummy.unsqueeze(0)).squeeze(0)
    # print(f'feat1: {feat1.shape}')
    feat1 = torch.cat([feat1, sti], dim=0).unsqueeze(0)
    input2 = torch.cat([dummy[1:], feat1], dim=0)
    # print(f'input2: {input2.shape}')
    feat2 = model(input2.unsqueeze(0)).squeeze(0)
    # print(f'feat2: {feat2.shape}')
    feat2 = torch.cat([feat2, sti], dim=0).unsqueeze(0)
    input3 = torch.cat([input2[1:], feat2], dim=0)
    # print(f'input3: {input3.shape}')
    feat3 = model(input3.unsqueeze(0)).squeeze(0)
    # print(f'feat3: {feat3.shape}')
    feat3 = torch.cat([feat3, sti], dim=0).unsqueeze(0)
    return torch.cat([dummy, feat1, feat2, feat3], dim=0)

def get_feat3_app(inputs):
    model_dir = inputs[0]
    data_dir = inputs[1]
    pth = glob.glob(f"{model_dir}/best_*.pth")
    if len(pth) == 0:
        raise ValueError(f"没有找到模型: {modelid}")
    pth = pth[0]
    mname = os.path.basename(pth).split(".")[0]
    temp = modelid.replace("lstm-2025", '')
    realid = f"{temp}-{mname}"
    json_path = data_dir + "/feat3/" + realid + ".json"
    if os.path.exists(json_path):
        print(f"skip {realid} because it exists")
        return ''
    model = try_get_model(pth)
    if model is None:
        return ''
    model.eval()
    model.to("cuda")
    feats = []
    with torch.no_grad():
        for each in get_dummy():
            feats.append([])
            for dummy in each:
                res = step3(model, dummy).cpu().numpy().tolist()
                feats[-1].append(res)
    feats_dict = {
        "model": modelid,
        "static_lefts": feats[0],
        "static_rights": feats[1],
        "move_lefts": feats[2],
        "move_rights": feats[3],
    }
    os.makedirs(DATAD + "/feat3", exist_ok=True)
    with open(json_path, "w") as f:
        json.dump(feats_dict, f)
    print(f"feat3 save to {json_path}")
    return json_path

def vis_feat3_app(inputs):
    json_path = inputs
    with open(json_path, "r") as f:
        feats_dict = json.load(f)
        sns.set_style("whitegrid")
        sns.set_context("notebook", font_scale=1.5)
        plt.rcParams['legend.fontsize'] = 14
        plt.rcParams['axes.labelsize'] = 16
        plt.rcParams['axes.titlesize'] = 18
        
        fig, axes = plt.subplots(2, 1, figsize=(15, 12))
        data_types = ['static_left', 'static_right', 'move_left', 'move_right']
        colors = ['#FF9999', '#66B2FF', '#99FF99', '#FFCC99']
        
        trajectory_data = []
        for i, data_type in enumerate(data_types):
            data = np.array(feats_dict[data_type + 's'][2]) # 30Hz
            for j in range(len(data)):
                trajectory_data.append({
                    'x': data[j, 0],
                    'y': data[j, 1],
                    'type': data_type,
                    'time': j
                })
        trajectory_df = pd.DataFrame(trajectory_data)
        
        sns.scatterplot(data=trajectory_df, x='x', y='y', hue='type', palette=dict(zip(data_types, colors)), 
                        ax=axes[0], s=50)
        for i, data_type in enumerate(data_types):
            data = trajectory_df[trajectory_df['type'] == data_type]
            x_coords = data['x'].values
            y_coords = data['y'].values
            axes[0].plot(x_coords, y_coords, color=colors[i], linewidth=2)
        
        axes[0].set_title('Trajectory Prediction')
        axes[0].set_xlabel('X Coordinate')
        axes[0].set_ylabel('Y Coordinate')
        axes[0].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        
        angle_data = []
        for i, data_type in enumerate(data_types):
            data = np.array(feats_dict[data_type+'s'][2])
            for j in range(len(data)):
                angle_data.append({
                    'angle': data[j, 2],
                    'type': data_type,
                    'time': j
                })
        angle_df = pd.DataFrame(angle_data)
        
        sns.scatterplot(data=angle_df, x='time', y='angle', hue='type', palette=dict(zip(data_types, colors)), 
                        ax=axes[1], s=50)
        for data_type in data_types:
            data = angle_df[angle_df['type'] == data_type]
            sns.lineplot(data=data, x='time', y='angle', color=colors[data_types.index(data_type)], 
                        ax=axes[1], linewidth=2)
        
        axes[1].set_title('Angle Prediction')
        axes[1].set_xlabel('Time Point')
        axes[1].set_ylabel('Angle')
        axes[1].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        
        plt.tight_layout()
        save_path = json_path.replace(".json", "_vis.png")
        plt.savefig(save_path, bbox_inches='tight', dpi=300)
        plt.close()
        
        print(f"Visualization saved to {save_path}")
        return save_path

def get_feat3(modelid):
    pth = glob.glob(f"{MODEL_DIR}/{modelid}/best_*.pth")
    if len(pth) == 0:
        print(f"没有找到模型: {modelid}")
        return
    pth = pth[0]
    mname = os.path.basename(pth).split(".")[0]
    temp = modelid.replace("lstm-2025", '')
    realid = f"{temp}-{mname}"
    json_path = DATAD + "/feat3/" + realid + ".json"
    if os.path.exists(json_path):
        print(f"skip {realid} because it exists")
        return
    model = try_get_model(pth)
    if model is None:
        return
    model.eval()
    model.to("cuda")
    feats = []
    with torch.no_grad():
        for each in get_dummy():
            feats.append([])
            for dummy in each:
                res = step3(model, dummy).cpu().numpy().tolist()
                feats[-1].append(res)
    feats_dict = {
        "model": modelid,
        "static_lefts": feats[0],
        "static_rights": feats[1],
        "move_lefts": feats[2],
        "move_rights": feats[3],
    }
    os.makedirs(DATAD + "/feat3", exist_ok=True)
    with open(json_path, "w") as f:
        json.dump(feats_dict, f)
    print(f"feat3 save to {json_path}")

def vis_feat3(modelid):
    json_path = datad + "/feat3/" + modelid + ".json"
    with open(json_path, "r") as f:
        feats_dict = json.load(f)
        sns.set_style("whitegrid")
        sns.set_context("notebook", font_scale=1.5)
        plt.rcParams['legend.fontsize'] = 14
        plt.rcParams['axes.labelsize'] = 16
        plt.rcParams['axes.titlesize'] = 18
        
        fig, axes = plt.subplots(2, 1, figsize=(15, 12))
        data_types = ['static_left', 'static_right', 'move_left', 'move_right']
        colors = ['#FF9999', '#66B2FF', '#99FF99', '#FFCC99']
        
        trajectory_data = []
        for i, data_type in enumerate(data_types):
            data = np.array(feats_dict[data_type + 's'][2])
            for j in range(len(data)):
                trajectory_data.append({
                    'x': data[j, 0],
                    'y': data[j, 1],
                    'type': data_type,
                    'time': j
                })
        trajectory_df = pd.DataFrame(trajectory_data)
        
        sns.scatterplot(data=trajectory_df, x='x', y='y', hue='type', palette=dict(zip(data_types, colors)), 
                        ax=axes[0], s=50)
        for i, data_type in enumerate(data_types):
            data = trajectory_df[trajectory_df['type'] == data_type]
            x_coords = data['x'].values
            y_coords = data['y'].values
            axes[0].plot(x_coords, y_coords, color=colors[i], linewidth=2)
        
        axes[0].set_title('Trajectory Prediction')
        axes[0].set_xlabel('X Coordinate')
        axes[0].set_ylabel('Y Coordinate')
        axes[0].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        
        angle_data = []
        for i, data_type in enumerate(data_types):
            data = np.array(feats_dict[data_type+'s'][2])
            for j in range(len(data)):
                angle_data.append({
                    'angle': data[j, 2],
                    'type': data_type,
                    'time': j
                })
        angle_df = pd.DataFrame(angle_data)
        
        sns.scatterplot(data=angle_df, x='time', y='angle', hue='type', palette=dict(zip(data_types, colors)), 
                        ax=axes[1], s=50)
        for data_type in data_types:
            data = angle_df[angle_df['type'] == data_type]
            sns.lineplot(data=data, x='time', y='angle', color=colors[data_types.index(data_type)], 
                        ax=axes[1], linewidth=2)
        
        axes[1].set_title('Angle Prediction')
        axes[1].set_xlabel('Time Point')
        axes[1].set_ylabel('Angle')
        axes[1].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        
        plt.tight_layout()
        save_path = datad + "/feat3/" + modelid + "_vis.png"
        plt.savefig(save_path, bbox_inches='tight', dpi=300)
        plt.close()
        
        print(f"Visualization saved to {save_path}")
        return save_path

def vis_feat3_direct(inputs):
    modelid = inputs[1]
    inputs = inputs[0]
    sns.set_style("whitegrid")
    sns.set_context("notebook", font_scale=1.5)
    plt.rcParams['legend.fontsize'] = 14
    plt.rcParams['axes.labelsize'] = 16
    plt.rcParams['axes.titlesize'] = 18
    fig, axes = plt.subplots(2, 1, figsize=(15, 12))
    data_types = ['static_left', 'static_right', 'move_left', 'move_right']
    colors = ['#FF9999', '#66B2FF', '#99FF99', '#FFCC99']
    trajectory_data = []
    for data, data_type in zip(inputs, data_types):
        data = np.array(data[2])
        for j in range(len(data)):
            trajectory_data.append({
                'x': data[j, 0],
                'y': data[j, 1],
                'type': data_type,
                'time': j
            })
        trajectory_df = pd.DataFrame(trajectory_data)
        sns.scatterplot(data=trajectory_df, x='x', y='y', hue='type', palette=dict(zip(data_types, colors)), 
                        ax=axes[0], s=50)
        for i, data_type in enumerate(data_types):
            data = trajectory_df[trajectory_df['type'] == data_type]
            x_coords = data['x'].values
            y_coords = data['y'].values
            axes[0].plot(x_coords, y_coords, color=colors[i], linewidth=2)
        axes[0].set_title('Trajectory Prediction')
        axes[0].set_xlabel('X Coordinate')
        axes[0].set_ylabel('Y Coordinate')
        axes[0].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        angle_data = []
        for data, data_type in zip(inputs, data_types):
            data = np.array(data[2])
            for j in range(len(data)):
                angle_data.append({
                    'angle': data[j, 2],
                    'type': data_type,
                    'time': j
                })
        angle_df = pd.DataFrame(angle_data)
        sns.scatterplot(data=angle_df, x='time', y='angle', hue='type', palette=dict(zip(data_types, colors)), 
                        ax=axes[1], s=50)
        for data_type in data_types:
            data = angle_df[angle_df['type'] == data_type]
            sns.lineplot(data=data, x='time', y='angle', color=colors[data_types.index(data_type)], 
                        ax=axes[1], linewidth=2)
        axes[1].set_title('Angle Prediction')
        axes[1].set_xlabel('Time Point')
        axes[1].set_ylabel('Angle')
        axes[1].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        plt.tight_layout()
        save_path = DATAD + "/feat3/" + modelid + "_vis.png"
        plt.savefig(save_path, bbox_inches='tight', dpi=300)
        plt.close()
        print(f"Visualization saved to {save_path}")
        return save_path

if __name__ == "__main__":
    datad = DATAD
    get_feat3("lstm-20250518164729")
    # for each in glob.glob(MODEL_DIR + "/lstm*"):
    #     each = os.path.basename(each)
    #     mid = each.replace('lstm-2025', '')
    #     if glob.glob(datad + "/feat3/" + mid + "*.json"):
    #         print(f"skip test {each}")
    #         continue
    #     get_feat3(each)
    # for each in glob.glob(datad + "/feat3/*.json"):
    #     each = os.path.basename(each)
    #     modelid = each.split(".")[0]
    #     if os.path.exists(datad + "/feat3/" + modelid + "_vis.png"):
    #         print(f"skip vis {modelid}")
    #         continue
    #     vis_feat3(modelid)
    # modelid = '0518164729-best_4231'
    # vis_feat3(modelid)
